import torch.nn as nn
import torch

from modules.Layer import *

class Decoder(nn.Module):
    def __init__(self, vocab, config):
        super(Decoder, self).__init__()
        self.vocab = vocab
        self.config = config

        self.cos = nn.CosineSimilarity(dim=1, eps=1e-08)

    def forward(self, pooled_s_x, pooled_t_x, pooled_neg_s, pooled_neg_t):
        s_loss = self.forward_cos_loss(pooled_s_x, pooled_t_x, pooled_neg_t)
        t_loss = self.forward_cos_loss(pooled_t_x, pooled_s_x, pooled_neg_s)
        return (s_loss + t_loss) / 2
    
    def forward_cos_loss(self, source, target, neg_targets):
        b, n, h = neg_targets.size()

        pos_cos = self.cos(source, target) 
        neg_cos = self.cos(
            source.unsqueeze(1).repeat(1, n, 1).view(b * n, h), 
            neg_targets.view(b * n, h)
        ) 

        neg_cos = neg_cos.view(b, n)

        all_loss = - torch.log(torch.exp(pos_cos) / (torch.exp(pos_cos) + torch.sum(torch.exp(neg_cos), dim=-1)) + 1e-5)

        loss = torch.mean(all_loss)
        return loss